{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "import segmentation_models_pytorch as smp\n",
    "from segmentation_models_pytorch.encoders.efficientnet import EfficientNetEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ToyModel(nn.Module):\n",
    "    \n",
    "    def __init__(self, dev0, dev1):\n",
    "        super(ToyModel, self).__init__()\n",
    "        self.dev0 = dev0\n",
    "        self.dev1 = dev1\n",
    "        self.net1 = torch.nn.Linear(10, 10).to(self.dev0)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.net2 = torch.nn.Linear(10, 5).to(self.dev1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.net1(x.to(self.dev0)))\n",
    "        return self.net2(x.to(self.dev1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ToyModel('cuda:0', 'cuda:1')\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.001)\n",
    "\n",
    "optimizer.zero_grad()\n",
    "outputs = model(torch.randn(20, 10))\n",
    "labels = torch.randn(20, 5).to('cuda:1')\n",
    "loss_fn(outputs, labels).backward()\n",
    "optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu_fn(x):\n",
    "    \"\"\" Swish activation function \"\"\"\n",
    "    return x * torch.sigmoid(x)\n",
    "\n",
    "class ParallelEfficientNetEncoder(EfficientNetEncoder):\n",
    "    \n",
    "    def __init__(self,\n",
    "        dev0, dev1,\n",
    "        width_coeff,\n",
    "        depth_coeff,\n",
    "        image_size,\n",
    "        dropout_rate,\n",
    "        drop_connect_rate,\n",
    "        block_chunks,\n",
    "        in_channels = 3  # rgb\n",
    "    ):\n",
    "        super().__init__(width_coeff, depth_coeff, image_size, dropout_rate,\n",
    "                         drop_connect_rate, block_chunks, in_channels)\n",
    "        self.dev0, self.dev1 = dev0, dev1\n",
    "        \n",
    "        # split layers across devices\n",
    "        self._bn0 = self._bn0.to(dev0)\n",
    "        self._conv_stem = self._conv_stem.to(dev0)\n",
    "        self.blocks_to_device(self.block_chunks[0], self.block_chunks[1], dev1)\n",
    "        self.blocks_to_device(self.block_chunks[1], self.block_chunks[2], dev0)\n",
    "        self.blocks_to_device(self.block_chunks[2], self.block_chunks[3], dev1)\n",
    "        self.blocks_to_device(self.block_chunks[3], self.block_chunks[4], dev0)\n",
    "        \n",
    "    def blocks_to_device(self, start_idx, end_idx, device):\n",
    "        for idx in range(start_idx, end_idx):\n",
    "            self._blocks[idx].to(device)\n",
    "            \n",
    "    def forward(self, x):\n",
    "        x0 = relu_fn(self._bn0(self._conv_stem(x.to(self.dev0))))\n",
    "        x1 = self.forward_blocks(x0.to(self.dev1), self.block_chunks[0], self.block_chunks[1])\n",
    "        x2 = self.forward_blocks(x1.to(self.dev0), self.block_chunks[1], self.block_chunks[2])\n",
    "        x3 = self.forward_blocks(x2.to(self.dev1), self.block_chunks[2], self.block_chunks[3])\n",
    "        x4 = self.forward_blocks(x3.to(self.dev0), self.block_chunks[3], self.block_chunks[4])\n",
    "        return [x4.to(self.dev1), x3, x2.to(self.dev1), x1, x0.to(self.dev1)]\n",
    "    \n",
    "class PipelineEfficientNetEncoder(EfficientNetEncoder):\n",
    "    \n",
    "    def __init__(self,\n",
    "        dev0, dev1, split_size,\n",
    "        width_coeff,\n",
    "        depth_coeff,\n",
    "        image_size,\n",
    "        dropout_rate,\n",
    "        drop_connect_rate,\n",
    "        block_chunks,\n",
    "        in_channels = 3  # rgb\n",
    "    ):\n",
    "        super().__init__(width_coeff, depth_coeff, image_size, dropout_rate,\n",
    "                         drop_connect_rate, block_chunks, in_channels)\n",
    "        self.dev0, self.dev1 = dev0, dev1\n",
    "        self.split_size = split_size\n",
    "        \n",
    "        # split layers across devices\n",
    "        self._bn0 = self._bn0.to(dev0)\n",
    "        self._conv_stem = self._conv_stem.to(dev0)\n",
    "        self.blocks_to_device(self.block_chunks[0], self.block_chunks[1], dev1)\n",
    "        self.blocks_to_device(self.block_chunks[1], self.block_chunks[2], dev0)\n",
    "        self.blocks_to_device(self.block_chunks[2], self.block_chunks[3], dev1)\n",
    "        self.blocks_to_device(self.block_chunks[3], self.block_chunks[4], dev0)\n",
    "        \n",
    "    def blocks_to_device(self, start_idx, end_idx, device):\n",
    "        for idx in range(start_idx, end_idx):\n",
    "            self._blocks[idx].to(device)\n",
    "            \n",
    "    def forward(self, x):\n",
    "        x0 = relu_fn(self._bn0(self._conv_stem(x.to(self.dev0))))\n",
    "        x1 = self.forward_blocks(x0.to(self.dev1), self.block_chunks[0], self.block_chunks[1])\n",
    "        x2 = self.forward_blocks(x1.to(self.dev0), self.block_chunks[1], self.block_chunks[2])\n",
    "        x3 = self.forward_blocks(x2.to(self.dev1), self.block_chunks[2], self.block_chunks[3])\n",
    "        x4 = self.forward_blocks(x3.to(self.dev0), self.block_chunks[3], self.block_chunks[4])\n",
    "        return [x4.to(self.dev1), x3, x2.to(self.dev1), x1, x0.to(self.dev1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[352, 120, 48, 24, 32]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ctor_kwargs = {\n",
    "    'dev0': 'cuda:0',\n",
    "    'dev1': 'cuda:1',\n",
    "    'width_coeff': 1.1,\n",
    "    'depth_coeff': 1.2,\n",
    "    'image_size': 260,\n",
    "    'dropout_rate': 0.3,\n",
    "    'drop_connect_rate': 0.2,\n",
    "    'block_chunks': [0, 3, 8, 16, 23]\n",
    "}\n",
    "\n",
    "model = ParallelEfficientNetEncoder(**ctor_kwargs)\n",
    "out_shapes = model._out_shapes\n",
    "out_shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ParallelEfficientNetEncoder'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.__class__.__name__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = nn.BCEWithLogitsLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 3, 256, 256]) torch.Size([16, 32, 128, 128]) torch.Size([16, 24, 64, 64]) torch.Size([16, 40, 32, 32]) torch.Size([16, 112, 16, 16]) torch.Size([16, 320, 8, 8])\n"
     ]
    }
   ],
   "source": [
    "B, C, H, W = 16, 3, 256, 256\n",
    "# dummy data\n",
    "data = torch.rand((B, C, H, W))\n",
    "t0 = torch.rand((B, out_shapes[4], H // 2, W // 2)).to('cuda:1')\n",
    "t1 = torch.rand((B, out_shapes[3], H // 4, W // 4)).to('cuda:1')\n",
    "t2 = torch.rand((B, out_shapes[2], H // 8, W // 8)).to('cuda:1')\n",
    "t3 = torch.rand((B, out_shapes[1], H // 16, W // 16)).to('cuda:1')\n",
    "t4 = torch.rand((B, out_shapes[0], H // 32, W // 32)).to('cuda:1')\n",
    "print(data.size(), t0.size(), t1.size(), t2.size(), t3.size(), t4.size())\n",
    "targets = [t4, t3, t2, t1, t0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "splits = iter(data.split(4, dim=0))\n",
    "?splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_next = next(splits)\n",
    "?s_next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(batches):\n",
    "    for _ in range(batches):\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(data)\n",
    "        for i in range(5):\n",
    "            loss_fn(outputs[i], targets[i]).backward(retain_graph=True)\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.87 s ± 11.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "train(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1.87 s ± 11.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load result: None\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Unet(\n",
       "  (encoder): EfficientNetEncoder(\n",
       "    (_conv_stem): Conv2dStaticSamePadding(\n",
       "      3, 32, kernel_size=(3, 3), stride=(2, 2), bias=False\n",
       "      (static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)\n",
       "    )\n",
       "    (_bn0): BatchNorm2d(32, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "    (_blocks): ModuleList(\n",
       "      (0): MBConvBlock(\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          32, 32, kernel_size=(3, 3), stride=[1, 1], groups=32, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(32, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          32, 8, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          8, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(16, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (1): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(96, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          96, 96, kernel_size=(3, 3), stride=[2, 2], groups=96, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(96, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          96, 4, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          4, 96, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(24, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (2): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          144, 144, kernel_size=(3, 3), stride=(1, 1), groups=144, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          144, 6, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          6, 144, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          144, 24, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(24, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (3): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          144, 144, kernel_size=(5, 5), stride=[2, 2], groups=144, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 2, 1, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          144, 6, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          6, 144, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          144, 40, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(40, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (4): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          240, 240, kernel_size=(5, 5), stride=(1, 1), groups=240, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          240, 10, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          10, 240, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          240, 40, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(40, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (5): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          240, 240, kernel_size=(3, 3), stride=[2, 2], groups=240, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          240, 10, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          10, 240, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          240, 80, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(80, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (6): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          480, 480, kernel_size=(3, 3), stride=(1, 1), groups=480, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          480, 20, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          20, 480, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          480, 80, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(80, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (7): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          480, 480, kernel_size=(3, 3), stride=(1, 1), groups=480, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          480, 20, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          20, 480, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          480, 80, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(80, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (8): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          480, 480, kernel_size=(5, 5), stride=[1, 1], groups=480, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          480, 20, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          20, 480, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          480, 112, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(112, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (9): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          672, 672, kernel_size=(5, 5), stride=(1, 1), groups=672, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          672, 28, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          28, 672, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          672, 112, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(112, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (10): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          672, 672, kernel_size=(5, 5), stride=(1, 1), groups=672, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          672, 28, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          28, 672, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          672, 112, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(112, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (11): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          672, 672, kernel_size=(5, 5), stride=[2, 2], groups=672, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 2, 1, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          672, 28, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          28, 672, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          672, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (12): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(5, 5), stride=(1, 1), groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (13): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(5, 5), stride=(1, 1), groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (14): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(5, 5), stride=(1, 1), groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (15): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(3, 3), stride=[1, 1], groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 320, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(320, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (decoder): UnetDecoder(\n",
       "    (layer1): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(432, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer2): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(296, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer3): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(152, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer4): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer5): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_conv): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (activation): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unet = smp.Unet('efficientnet-b0')\n",
    "unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
